#ifndef SEGMENTATION_H
#define SEGMENTATION_H

#include <iostream>
#include <string>
#include <opencv2/opencv.hpp>
#include <unistd.h>
#include <omp.h>

#include "rt_model_operation.h"

using namespace algocommon;

namespace segmentation
{


    struct InitParams : public OnnxModeltoTRTParam
    {
        int gpu_id;
        std::string onnx_model_path;
        std::string trt_model_dir;
        std::string trt_model_name;
    };

    class Segment : public OnnxModeltoTRT
    {

    public:

        Segment(const InitParams &params);
        virtual ~Segment(){};

    public:

        cv::Mat Execute(const cv::Mat &img);

    private:

        Result InitNetwork();
        void PreProcessCPU(const cv::Mat &img);
        cv::Mat PostProcessCPU(const std::vector<TensorFloat>&output_vector);
        cv::Mat PostOutput(const std::vector<cv::Mat>&output_vector);
        void SetCutSize(const cv::Mat &img);

    private:

        InitParams params_;
        std::vector<float> mean_{0.485, 0.456, 0.406};
        std::vector<float> std_{0.229, 0.224, 0.225};
        int output_height_ = 512;
        int output_width_ = 512;
        int longsize_ =512;
        int ow_; //cut_image
        int oh_;
    };







}//namespace segmentation

#endif